import netCDF4 
import sys, struct
import math as pm
import numpy as np
import pylab as p
import math
import csv
import matplotlib.pyplot as plt
import matplotlib as mlt
import h5py
from netcdf_reader import *
from scipy import ndimage
import scipy.sparse.linalg

#OUTPUT_DATA_DIR = '/home/behollis/Dropbox/visweek2015/oslines/'
#INPUT_DATA_DIR = '/home/behollis/DATA/in/ts00050/'
OUTPUT_DATA_DIR = '/home/brad/visweek2015-revision-code/oslines/'
INPUT_DATA_DIR = '/home/brad/DATA/ocean/'


FILE_NAME = 'pe_dif_sep2_98.nc' 
FILE_NAME_CENTRAL_FORECAST = 'pe_fct_aug25_sep2.nc'
#INPUT_DATA_DIR = '/home/behollis/DATA/pierre/ocean/'
#OUTPUT_DATA_DIR = '/home/behollis/Dropbox/visweek2015/oslines/'
#realizations file 
pe_dif_sep2_98_file = INPUT_DATA_DIR + FILE_NAME
pe_fct_aug25_sep2_file = INPUT_DATA_DIR + FILE_NAME_CENTRAL_FORECAST 

COM =  2
LON = 53
LAT = 90
LEV = 16
DEPTH = -2.0

def showVelMag(vclin):
    mag = np.zeros(shape=(LAT,LON))
    for x in range(0,LAT):
        for y in range(0,LON):
            avg_vel_mag = 0.0
            for idx in range(7,8):
                avg_vel_mag += math.sqrt( math.pow(vclin[idx][x,y,0][0],2) + math.pow(vclin[idx][x,y,0][1],2) ) 
                
            mag[x,y] = avg_vel_mag# / float(TOTMEMS)
    
   
    plt.imshow(mag, origin='lower', interpolation=None, vmax = 28, cmap='gist_ncar')
    plt.colorbar()
    
    plt.imshow(landmask, origin='lower', interpolation=None)
    
    '''
    mems = [ d for d in f ]
    for m in mems:
        x = '/x020'
        y = '/y120'
        sl = f[m+x+y]
        plt.plot(sl[0], sl[1], color ='white')
        print len(sl[0])
        
    plt.plot(20., 120., '+', color='red', linewidth=50.0, ms=10.0) # draw seed loc
    '''

    plt.show()
    
    f.close()


def showSlines(file):
    
    plt.imshow(landmask, origin='lower', interpolation=None)
        
    f = h5py.File(OUTPUT_DATA_DIR + file, 'r')
    mems = [ d for d in f ]
    
    for m in mems:
        x = '/y021'
        y = '/x054'
        sl = f[m+y+x]
        plt.plot(sl[1], sl[0])
        
    plt.plot(21., 54., '+', color='red', linewidth=10.0, ms=10.0) # draw seed loc

    
    plt.show()
    f.close()

def showFTVA(file):
    
    f = h5py.File(OUTPUT_DATA_DIR + file, 'r')
   
    ftva = np.zeros(shape=(LAT,LON))
   
    for x in range(0,LAT):
        for y in range(0, LON):
            c = f['cov'][x,y]
            
            w, v = np.linalg.eig(c)
            sorted(w)
                
            ftva[x,y] = w[0]
    
    plt.imshow(ftva, origin='lower', interpolation=None)
    plt.imshow(landmask, origin='lower', interpolation=None)
    plt.colorbar()
    plt.show()
    f.close()
    
def makeLandMask():
    rreader = NetcdfReader(pe_dif_sep2_98_file)
    landv = rreader.readVarArray('landv')
    
    landmask = np.zeros(shape=(LAT,LON,4))
    for x in range(0,LAT):
        for y in range(0, LON):
            if landv[x,y] == 0.0:
                landmask[x,y] = (0.2,0.2,0.2,1)
            else:
                landmask[x,y] = (0,0,0,0)
                
    return landmask

def showEntropySlineClusters(file,x=None,y=None):
    f = h5py.File(OUTPUT_DATA_DIR + file, 'r')
    termClusters = f['slineclusters'] 
    
    max = int(np.amax(termClusters))
    
    FACT = 8
    
    # define the colormap
    cmap = plt.cm.jet
    # extract all colors from the .jet map
    cmaplist = [cmap(i) for i in range(0,max*FACT*(max+1),max*FACT)]
    # force the first color entry to be grey
    cmaplist[0] = landmask[0,0]#(.5,.5,.5,1.0)
    # create the new map
    cmap = cmap.from_list('Custom cmap', cmaplist, len(cmaplist))
    
    # define the bins and normalize
    bounds = np.linspace(0,max+1,max+2)
    norm = mlt.colors.BoundaryNorm(bounds, cmap.N)
    
    plt.imshow(termClusters, origin='lower',cmap=cmap, norm=norm, interpolation=None)
    plt.colorbar()
    
    #plot slines at seed
    if x is not None:
        f = h5py.File(OUTPUT_DATA_DIR + 'oceanSlinesLevel00.hdf5', 'r')
        mems = [ d for d in f ]
        
        for m in mems:
            dir = m + '/x' + str(y).zfill(3) + '/y' + str(x).zfill(3)
            sl = f[dir]
            plt.plot( sl[1], sl[0] )
            
        plt.plot(x, y, '+', color='white', linewidth=10.0, ms=10.0) # draw seed loc
    
    
    plt.imshow(landmask, origin='lower', interpolation=None)
    plt.show()
    f.close()

def showTerminalParticleClusters(file):
    f = h5py.File(OUTPUT_DATA_DIR + file, 'r')
    termClusters = f['ptclusters'] 
    
    max = int(np.amax(termClusters))
    
    # define the colormap
    cmap = plt.cm.jet
    # extract all colors from the .jet map
    cmaplist = [cmap(i) for i in range(0,max*10*(max+1),max*10)]
    # force the first color entry to be grey
    cmaplist[0] = landmask[0,0]#(.5,.5,.5,1.0)
    # create the new map
    cmap = cmap.from_list('Custom cmap', cmaplist, len(cmaplist))
    
    # define the bins and normalize
    bounds = np.linspace(0,max+1,max+2)
    norm = mlt.colors.BoundaryNorm(bounds, cmap.N)
    
    '''
    cb = mlt.colorbar.ColorbarBase(ax2, cmap=cmap, norm=norm, spacing='proportional', \
                                   ticks=bounds, boundaries=bounds, format='%1i')
    '''
     
    plt.imshow(termClusters, origin='lower',cmap=cmap, norm=norm, interpolation=None)
    plt.colorbar()
    plt.imshow(landmask, origin='lower', interpolation=None)
    plt.show()
    f.close()
    
def showEntropyFields(file):
    f = h5py.File(OUTPUT_DATA_DIR + file, 'r')
    
    avgAngEntr = f['avgAngularEntropy']
    avgAngEntrMin = np.min(avgAngEntr)
    avgAngEntrMax = np.max(avgAngEntr)
    avgLinEntr = f['avgLinearEntropy']
    avgLinEntrMin = np.min(avgLinEntr)
    avgLinEntrMax = np.max(avgLinEntr)
    
    
    '''
    dsampled = output.create_dataset(name='ptsSampledForClustering', shape=(LAT,LON), dtype='f')
    dentroL = output.create_dataset(name='avgAngularEntropy', shape=(LAT,LON), dtype='f')
    dentroA = output.create_dataset(name='avgLinearEntropy', shape=(LAT,LON), dtype='f')
    '''

    
#    plt.imshow(avgAngEntr, origin='lower', interpolation=None, vmin=0, vmax=0.5)
    plt.imshow(avgAngEntr, origin='lower', interpolation=None, vmin=avgAngEntrMin, vmax=avgAngEntrMax)
    cb = plt.colorbar(); 
    plt.imshow(landmask, origin='lower', interpolation=None)
    plt.show()
    
#    plt.imshow(avgLinEntr, origin='lower', interpolation=None, vmin=0, vmax=4)
    plt.imshow(avgLinEntr, origin='lower', interpolation=None, vmin=avgLinEntrMax, vmax=avgLinEntrMax)
    plt.colorbar(); 
    plt.imshow(landmask, origin='lower', interpolation=None)
    plt.show()
    
    f.close()
    
    
def loadUVFields():
    
    SEED_LEVEL = 0
    level = SEED_LEVEL
    
    #realizations file 
    pe_dif_sep2_98_file = INPUT_DATA_DIR + FILE_NAME
    pe_fct_aug25_sep2_file = INPUT_DATA_DIR + FILE_NAME_CENTRAL_FORECAST 
    
    #realizations reader 
    rreader = NetcdfReader(pe_dif_sep2_98_file)
    
    #central forecasts reader 
    creader = NetcdfReader(pe_fct_aug25_sep2_file)
    vclin8 = creader.readVarArray('vclin', 7)
    
    #deviations from central forecast for all 600 realizations
    vclin = rreader.readVarArray('vclin')  
    vclin = addCentralForecast(vclin, vclin8, level_start=SEED_LEVEL, level_end=SEED_LEVEL)
    
    return vclin

def showEntropyFields(file):
    f = h5py.File(OUTPUT_DATA_DIR + file, 'r')
    
    avgAngEntr = f['avgAngularEntropy']
    avgAngEntrMin = np.min(avgAngEntr)
    avgAngEntrMax = np.max(avgAngEntr)
    avgLinEntr = f['avgLinearEntropy']
    avgLinEntrMin = np.min(avgLinEntr)
    avgLinEntrMax = np.max(avgLinEntr)
    
    
    # angular entropy map
    cmap = plt.cm.jet
    
    '''
    dsampled = output.create_dataset(name='ptsSampledForClustering', shape=(LAT,LON), dtype='f')
    dentroL = output.create_dataset(name='avgAngularEntropy', shape=(LAT,LON), dtype='f')
    dentroA = output.create_dataset(name='avgLinearEntropy', shape=(LAT,LON), dtype='f')
    '''
    #    plt.imshow(avgAngEntr, origin='lower', interpolation=None, vmin=0, vmax=0.5)
    plt.imshow(avgAngEntr, origin='lower', interpolation=None, vmin=0.8, vmax=avgAngEntrMax)
    cb = plt.colorbar(); 
    plt.imshow(landmask, origin='lower', interpolation=None)
    plt.show()
    
#    plt.imshow(avgLinEntr, origin='lower', interpolation=None, vmin=0, vmax=4)
    plt.imshow(avgLinEntr, origin='lower', interpolation=None, vmin=0.8, vmax=1.0)
    plt.colorbar(); 
    plt.imshow(landmask, origin='lower', interpolation=None)
    plt.show()
    
    '''
    # sum of maps
    plt.imshow(np.add(avgLinEntr,avgAngEntr), origin='lower', interpolation=None, vmin=0.6, vmax=1.0, cmap=cmap)
    plt.colorbar(); 
    plt.show()
    '''
    
    powOf2 = np.zeros(shape=avgLinEntr.shape); powOf2.fill(2.0)
    sqRoot = np.zeros(shape=avgLinEntr.shape); sqRoot.fill(0.5)
    
    gradLinEntr = np.gradient(avgLinEntr)
    gradLinEntrMag = np.power( np.power(gradLinEntr[0], powOf2) + np.power(gradLinEntr[1], powOf2), sqRoot)
    
    gradAngEntr = np.gradient(avgAngEntr)
    gradAngEntrMag = np.power( np.power(gradAngEntr[0], powOf2) + np.power(gradAngEntr[1], powOf2), sqRoot)
    
    # GRADIENT of LINEAR ENTROPY map
    plt.imshow(gradAngEntrMag, origin='lower', interpolation=None, vmin=np.min(gradAngEntrMag), vmax=0.05, cmap=cmap)
    plt.colorbar();
    plt.imshow(landmask, origin='lower', interpolation=None) 
    plt.show()
    
    # GRADIENT of A ENTROPY map
    plt.imshow(gradLinEntrMag, origin='lower', interpolation=None, vmin=np.min(gradLinEntrMag), vmax=0.05, cmap=cmap)
    plt.colorbar(); 
    plt.imshow(landmask, origin='lower', interpolation=None)
    plt.show()
    
    f.close()  
    
def showSlineSamplingRate(file, mult=10):
    f = h5py.File(OUTPUT_DATA_DIR + file, 'r')
    termClusters = f['ptsSampledForClustering'] 
    
    max = int(np.amax(termClusters))
    
    # define the colormap
    cmap = plt.cm.jet
    # extract all colors from the .jet map
    cmaplist = [cmap(i) for i in range(0,max*mult*(max+1),max*mult)]
    # force the first color entry to be grey
    # create the new map
    cmap = cmap.from_list('Custom cmap', cmaplist, len(cmaplist))
    
    # define the bins and normalize
    bounds = np.linspace(0,max+1,max+2)
    norm = mlt.colors.BoundaryNorm(bounds, cmap.N)
    
    imgplot = plt.imshow(termClusters, origin='lower',cmap=cmap, norm=norm, interpolation=None)
    
    '''
    ps = list()
    
    for idx in range( 0,len(cmaplist) ):
        p, = mpatches.Patch(color='red', label='2' )
        ps.append(p) 
        
    plt.legend(handles=ps)
    '''

    plt.colorbar()
    plt.imshow(landmask, origin='lower', interpolation=None) 
    
    plt.show()
    f.close()
      
    
if __name__ == '__main__':    
    landmask = makeLandMask()
#    vclin = loadUVFields()
#    showVelMag(vclin)
    
    
    
#    showSlines(file='oceanSlinesLevel00.hdf5')
#    showFTVA(file='oceanFtvaBackwardLev08.hdf5')
#    showTerminalParticleClusters(file='oceanTerminalClusterBackwardLev08.hdf5')
#    showEntropySlineClusters(file='ocean.SlineClusters.varSampling.3to13pts.hdf5')
#    showEntropySlineClusters(file='oceanEntropySlineClustersLev00.8pts.hdf5')

    showEntropySlineClusters(file='ocean.SlineClusters.varSampling.3to13pts.hdf5')
    showSlineSamplingRate(file='ocean.SlineClusters.varSampling.3to13pts.hdf5', mult=1)

#    showEntropyFields(file='oceanEntropiesONLY.fixed.hdf5')
    
    
    
    
    
    